from fastapi import Request
from fastuuid import fastuuid
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint

from app.constants import HEADER_TRACE_ID
from app.core.log import logger
from app.utils import context


class TraceMiddleware(BaseHTTPMiddleware):
    """
    链路跟踪 中间件

    用于实现分布式系统的链路追踪功能。通过在请求头中传递trace_id，
    可以追踪一个请求在多个服务间的调用链路，便于问题排查和性能分析。
    """

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        """
        中间件主入口函数，处理每个 HTTP 请求的链路跟踪逻辑

        :param request: FastAPI 的 Request 对象，包含客户端请求信息
        :param call_next: 下一个中间件或路由处理函数
        :return: 返回处理后的响应结果
        """
        logger.debug("----------> 2. 链路跟踪 中间件")

        # 从请求头中获取trace_id
        trace_id = request.headers.get(HEADER_TRACE_ID)

        # 如果trace_id不存在，则生成一个
        if trace_id is None:
            logger.debug("请求头中的trace_id不存在, 生成一个")
            trace_id = str(fastuuid.uuid4()).replace("-", "")

        # 设置到Mdc上下文中
        context.set_trace_id(trace_id)

        try:
            # 继续执行下一个中间件或路由处理器
            response = await call_next(request)

            logger.debug("<---------- 2. 链路跟踪 中间件")
            return response
        finally:
            # 请求处理完成后，统一清理上下文环境
            # 确保即使发生异常也能正确清理资源
            context.remove()
